import os

from langchain.agents import AgentExecutor
from langchain.chat_models import QianfanChatEndpoint
from langchain.tools import WikipediaQueryRun, ArxivQueryRun
from langchain.utilities import ArxivAPIWrapper, WikipediaAPIWrapper
from qianfan.extensions.langchain.agents import QianfanSingleActionAgent

os.environ["QIANFAN_AK"] = "ZGgmCMM95MfabqYLG2swwRWM"
os.environ["QIANFAN_SK"] = "yDAIOpvEu7vm2phjYDT5sdIg5h4LBb4v"

qianfan_chat_model = QianfanChatEndpoint(model="ERNIE-Bot")

tools = [
    ArxivQueryRun(
        description="这是一个论文搜索的工具"
                    "当你需要查询论文的时候，请使用这个工具"
                    "输入应该是一条你所需要查询的论文的信息"
        ,
        api_wrapper=ArxivAPIWrapper(top_k_results=1)
    ),
    WikipediaQueryRun(
        description="这是一个维基百科的调用工具."
                    "当你需要查询有关论文作者、历史的百科信息时，请使用这个工具"
                    "输入应该是一条你所需要查询的内容",
        api_wrapper=WikipediaAPIWrapper(load_all_available_meta=True)
    )
]

qianfan_agent = QianfanSingleActionAgent.from_system_prompt(tools, qianfan_chat_model)
executor = AgentExecutor(agent=qianfan_agent, tools=tools, verbose=True)

result = executor.run("帮我查询physics领域论文，以及作者的百科信息")
print(result)
